from skimage import io
import numpy as np
import matplotlib.pyplot as plt
from skimage.color import rgb2gray, rgb2hsv, hsv2rgb
from skimage import transform
from random import shuffle
import random
import sys
import os

def background(rownum, colnum):
    unit = 167
    nulls= os.listdir(path='./images')

    num = rownum * colnum
    grids_num = 0
    grids = []
    while (grids_num < num):
        shuffle(nulls)
        grids += nulls.copy()
        grids_num += len(nulls)
    grids = grids[:num]

    black = np.zeros((unit * rownum , unit * colnum, 3),dtype=np.uint8)

    i = 0
    angles = [0, 0, 0, 90, 180, 180, 180, 270]
    angles = [0, 180]
    
    changes = [0, 0, 0, 0, 0, 1, 2, 3]
    for col in range(colnum):
        for row in range(rownum):
            name = grids[i]
            i += 1
            path = "./images/{}".format(name)
            im = io.imread(path)
            c = random.choice(changes)
            if c == 1:
                im[:, :, :] = im[::-1, :, :]
            elif c == 2:
                im[:, :, :] = im[:, ::-1, :]
            elif c == 3:
                im = transform.rotate(im, random.choice(angles), preserve_range=True)     

            black[unit*row:unit*(row+1),unit*col:unit*(col+1),:] = im[:,:,:]
    return black

def adjuct(black, rownum, colnum):
    u = 167
    b = rgb2hsv(black)
    vs = [1.1, 0.8] * ((rownum + 1) // 2)
    vs = vs[:rownum]
    vs = [x ** 1.1 for x in vs]
    for i, v in enumerate(vs):
        b[:, u*i:u*(i+1), 2] *= v

    for j in range(colnum):
        total_mean = b[:, j*u:(j+1)*u, 2].mean()
        grids_mean = []
        for i in range(rownum):
            m = b[i*u:(i+1)*u, j*u:(j+1)*u, 2].mean()
            grids_mean.append(m)

        for i in range(rownum):
            b[i*u:(i+1)*u, j*u:(j+1)*u, 2] *= total_mean / grids_mean[i]

    b = hsv2rgb(b)
    return b

def blackwhite(name):
    original = io.imread(name)

    grayscale = rgb2gray(original)

    ts = 87 / 255 
    mask = grayscale > ts
    grayscale[mask] = 1.0
    mask = grayscale <= ts
    grayscale[mask] = 0.0
    grayscale = 1.0 - grayscale
    return grayscale



def main(filename, rownum, colnum):
    black = background(rownum, colnum)
    black = adjuct(black, rownum, colnum)
    black = rgb2gray(black)

    chars = blackwhite(filename)

    black = transform.resize(black, chars.shape)

    mask_b = chars < 0.5
    mask_w = chars >= 0.5
   
    r = np.zeros((chars.shape), dtype=np.uint8)
    
    r[:] = chars[:]*255
    
    r[mask_b] = black[mask_b] * 255
    r[mask_w] = r[mask_w] / 2 + (black[mask_w]* -1.0 + 1.0) * 120
 
    plt.figure()
    plt.imshow(r, cmap=plt.cm.gray)
    plt.show()

    if filename == "input.jpg":
        filename = "output.jpg"
    else:
        p, f = os.path.split(filename)
        filename = "{}out_{}".format(p, f)
    io.imsave(filename, r)

if __name__ == "__main__":
    rownum = 6
    colnum = 4
    if len(sys.argv) == 1:
        filename = "input.jpg"
    elif len(sys.argv) == 2:
        filename = sys.argv[1]
    elif len(sys.argv) == 4:
        filename = sys.argv[1]
        rownum = int(sys.argv[2])
        colnum = int(sys.argv[3])
    else:
        print("usage: python3 blackbg.py <input.jpg> <rownum> <colnum>")
        sys.exit()
    
    main(filename, rownum, colnum)


